package com.gitee.zycra.file.utils;

import org.bouncycastle.asn1.gm.GMObjectIdentifiers;
import org.bouncycastle.crypto.CipherParameters;
import org.bouncycastle.crypto.engines.SM2Engine;
import org.bouncycastle.crypto.params.ECDomainParameters;
import org.bouncycastle.crypto.params.ECPrivateKeyParameters;
import org.bouncycastle.crypto.params.ECPublicKeyParameters;
import org.bouncycastle.crypto.params.ParametersWithRandom;
import org.bouncycastle.jcajce.provider.asymmetric.ec.BCECPrivateKey;
import org.bouncycastle.jcajce.provider.asymmetric.ec.BCECPublicKey;
import org.bouncycastle.jce.interfaces.ECKey;
import org.bouncycastle.jce.provider.BouncyCastleProvider;
import org.bouncycastle.jce.spec.ECParameterSpec;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.nio.charset.StandardCharsets;
import java.security.KeyFactory;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.SecureRandom;
import java.security.Signature;
import java.security.spec.ECGenParameterSpec;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.Base64;

/**
 * Encrypt or decrypt util of SM2.
 *
 * @author zycra
 * @since 1.0.0
 */
public final class SM2Util extends BCUtil {

    private static final Logger LOGGER = LoggerFactory.getLogger(SM2Util.class);

    /**
     * SM2 std name.
     *
     * @since 1.0.0
     */
    private static final String SM2_STD_NAME = "sm2p256v1";

    /**
     * ALGO name.
     *
     * @since 1.0.0
     */
    private static final String ALGO_NAME_EC = "EC";

    private SM2Util() {
    }

    /**
     * Create a new SM2 key pair.
     * <p>Use {@link #getPublicKey(KeyPair)} and {@link #getPrivateKey(KeyPair)} to get keys.
     *
     * @return key pair.
     * @since 1.0.0
     */
    public static KeyPair createECKeyPair() {
        try {
            ECGenParameterSpec sm2Spec = new ECGenParameterSpec(SM2_STD_NAME);
            KeyPairGenerator kpg = KeyPairGenerator.getInstance(ALGO_NAME_EC, BouncyCastleProvider.PROVIDER_NAME);
            kpg.initialize(sm2Spec, new SecureRandom());
            return kpg.generateKeyPair();
        } catch (Exception e) {
            LOGGER.error("create key failed", e);
        }
        return null;
    }

    /**
     * Get public key from the given key pair.
     *
     * @param keyPair key pair.
     * @return public key with base64 encode.
     * @since 1.0.0
     */
    public static String getPublicKey(KeyPair keyPair) {
        return Base64.getEncoder().encodeToString(keyPair.getPublic().getEncoded());
    }

    /**
     * Get private key from the given key pair.
     *
     * @param keyPair key pair.
     * @return private key with base64 encode.
     * @since 1.0.0
     */
    public static String getPrivateKey(KeyPair keyPair) {
        return Base64.getEncoder().encodeToString(keyPair.getPrivate().getEncoded());
    }

    /**
     * SM2 encrypt content by public key.
     *
     * @param content   the content to encrypt.
     * @param publicKey SM2 public key.
     * @return encrypt result.
     * @since 1.0.0
     */
    public static String encryptByPublicKey(String content, String publicKey) {
        return encryptByPublicKey(content, publicKey, SM2Engine.Mode.C1C3C2);
    }

    /**
     * SM2 encrypt content by public key.
     *
     * @param content   the content to encrypt.
     * @param publicKey SM2 public key.
     * @param mode      SM2 encrypt mode.
     * @return encrypt result.
     * @since 1.0.0
     */
    public static String encryptByPublicKey(String content, String publicKey, SM2Engine.Mode mode) {
        return doSM2(content, publicKey, mode, true);
    }

    /**
     * SM2 decrypt content by private key.
     *
     * @param content    the content to decrypt.
     * @param privateKey SM2 private key.
     * @return decrypt result.
     * @since 1.0.0
     */
    public static String decryptByPrivateKey(String content, String privateKey) {
        return decryptByPrivateKey(content, privateKey, SM2Engine.Mode.C1C3C2);
    }

    /**
     * SM2 decrypt content by private key.
     *
     * @param content    the content to decrypt.
     * @param privateKey SM2 private key.
     * @param mode       SM2 decrypt mode.
     * @return decrypt result.
     * @since 1.0.0
     */
    public static String decryptByPrivateKey(String content, String privateKey, SM2Engine.Mode mode) {
        return doSM2(content, privateKey, mode, false);
    }

    /**
     * Do SM2 encrypt or decrypt.
     *
     * @param content       the content to encrypt or decrypt.
     * @param key           public key for encrypt and private key for decrypt.
     * @param mode          SM2 encrypt or decrypt mode.
     * @param forEncryption true is to encrypt false is to decrypt.
     * @return encrypt or decrypt result.
     * @since 1.0.0
     */
    private static String doSM2(String content, String key, SM2Engine.Mode mode, boolean forEncryption) {
        if (content == null) {
            return null;
        }
        ECKey ecKey = forEncryption ? getPublicKey(key) : getPrivateKey(key);
        if (ecKey == null) {
            LOGGER.error("get key failed");
            return null;
        }
        ECParameterSpec ecParameterSpec = ecKey.getParameters();
        ECDomainParameters ecDomainParameters = new ECDomainParameters(ecParameterSpec.getCurve(), ecParameterSpec.getG(), ecParameterSpec.getN());
        SM2Engine sm2Engine = new SM2Engine(mode);
        byte[] data;
        CipherParameters cipherParameters;
        if (forEncryption) {
            data = content.getBytes(StandardCharsets.UTF_8);
            BCECPublicKey publicKey = (BCECPublicKey) ecKey;
            ECPublicKeyParameters ecPublicKeyParameters = new ECPublicKeyParameters(publicKey.getQ(), ecDomainParameters);
            cipherParameters = new ParametersWithRandom(ecPublicKeyParameters, new SecureRandom());
        } else {
            data = Base64.getMimeDecoder().decode(content);
            BCECPrivateKey privateKey = (BCECPrivateKey) ecKey;
            cipherParameters = new ECPrivateKeyParameters(privateKey.getD(), ecDomainParameters);
        }
        sm2Engine.init(forEncryption, cipherParameters);
        try {
            byte[] result = sm2Engine.processBlock(data, 0, data.length);
            return forEncryption ? Base64.getEncoder().encodeToString(result) : new String(result, StandardCharsets.UTF_8);
        } catch (Exception e) {
            LOGGER.error("do SM2 error", e);
        }
        return null;
    }

    /**
     * SM2 sign by private key.
     *
     * @param content    the content to sign.
     * @param privateKey SM2 private key.
     * @return sign result.
     * @since 1.0.0
     */
    public static String signByPrivateKey(String content, String privateKey) {
        if (content == null) {
            return null;
        }
        PrivateKey priKey = getPrivateKey(privateKey);
        if (priKey == null) {
            return null;
        }
        try {
            Signature signature = Signature.getInstance(GMObjectIdentifiers.sm2sign_with_sm3.toString(), BouncyCastleProvider.PROVIDER_NAME);
            signature.initSign(priKey);
            signature.update(content.getBytes(StandardCharsets.UTF_8));
            byte[] sign = signature.sign();
            return Base64.getEncoder().encodeToString(sign);
        } catch (Exception e) {
            LOGGER.error("do sign by private key error", e);
        }
        return null;
    }

    /**
     * SM2 check sign by public key.
     *
     * @param content   the content to check sign.
     * @param publicKey SM2 public key.
     * @param signStr   the sign string to check.
     * @return check result.
     * @since 1.0.0
     */
    public static boolean verifyByPublicKey(String content, String publicKey, String signStr) {
        if (content == null) {
            return false;
        }
        PublicKey pubKey = getPublicKey(publicKey);
        if (pubKey == null) {
            return false;
        }
        try {
            Signature signature = Signature.getInstance(GMObjectIdentifiers.sm2sign_with_sm3.toString(), BouncyCastleProvider.PROVIDER_NAME);
            signature.initVerify(pubKey);
            signature.update(content.getBytes(StandardCharsets.UTF_8));
            byte[] sign = Base64.getMimeDecoder().decode(signStr);
            return signature.verify(sign);
        } catch (Exception e) {
            LOGGER.error("do verify sign by public key error", e);
        }
        return false;
    }

    /**
     * Get SM2 public key by base64 key string.
     *
     * @param key base64 key string.
     * @return public key object.
     * @since 1.0.0
     */
    private static BCECPublicKey getPublicKey(String key) {
        if (key == null) {
            return null;
        }
        try {
            byte[] keyBytes = Base64.getMimeDecoder().decode(key);
            return (BCECPublicKey) KeyFactory.getInstance(ALGO_NAME_EC, BouncyCastleProvider.PROVIDER_NAME)
                    .generatePublic(new X509EncodedKeySpec(keyBytes));
        } catch (Exception e) {
            LOGGER.error("get public key error", e);
        }
        return null;
    }

    /**
     * Get SM2 private key by base64 key string.
     *
     * @param key base64 key string.
     * @return private key object.
     * @since 1.0.0
     */
    private static BCECPrivateKey getPrivateKey(String key) {
        if (key == null) {
            return null;
        }
        try {
            byte[] keyBytes = Base64.getMimeDecoder().decode(key);
            return (BCECPrivateKey) KeyFactory.getInstance(ALGO_NAME_EC, BouncyCastleProvider.PROVIDER_NAME)
                    .generatePrivate(new PKCS8EncodedKeySpec(keyBytes));
        } catch (Exception e) {
            LOGGER.error("get private key error", e);
        }
        return null;
    }
}
